import time
import logging
import psycopg2
import backtrader as bt
from datetime import datetime
import psycopg2.extras
import pandas as pd

from Units import QUnits, TdxSecurity, StockAnalysis
from apscheduler.schedulers.blocking import BlockingScheduler
logging.basicConfig(format="%(asctime)s %(filename)s(line:%(lineno)d) [%(levelname)s] : %(message)s", datefmt="%Y-%M-%d %H:%M:%S", level=logging.DEBUG)

class StockSchedule:
    def __init__(self, parent=[]):
        self.scheduler = BlockingScheduler()
        self.security = TdxSecurity()
        self.analysis = StockAnalysis()
        self.cerebro = bt.Cerebro()
        self.cerebro.addstrategy(TradeStrategy)  # 增加一个策略
        self.cerebro.broker.setcash(100000.0)
        # 设置固定滑点
        self.cerebro.broker.set_slippage_fixed(fixed=0.001)
        # 设置 0.00005 = 0.005% 的手续费
        self.cerebro.broker.setcommission(commission=0.00005)
        self.df5m = dict()
        self.df1d = dict()
        self.stockList = parent
        for row in parent:
            data5m = self._read_data_5m(row['ths_code'])
            data1d = self._read_data_1d(row['ths_code'])
            now5m = self.security.getStockBarsFromApi(0, row['market'], row['stock_code'], 1, 300, True)
            now1d = self.security.getStockBarsFromApi(4, row['market'], row['stock_code'], 1, 50, True)
            self.df5m[row['tdx_code']] = {'data': data5m[0], 'count': data5m[1], 'start': data5m[2], 'end': data5m[3]}
            self.df1d[row['tdx_code']] = {'data': data1d[0], 'count': data1d[1], 'start': data1d[2], 'end': data1d[3]}

    def job_executes(self, rule='5m'):
        logging.info("开始执行任务，任务参数：%s" % rule)
        for r in self.stockList:
            data5m = self.df5m[r['tdx_code']]['data']
            data1d = self.df1d[r['tdx_code']]['data']
            now5m = self.security.getStockBarsFromApi(0, r['market'], r['stock_code'], 1, 150, True)
            now1d = self.security.getStockBarsFromApi(4, r['market'], r['stock_code'], 1, 30, True)
            data_5m = pd.concat([data5m, now5m], ignore_index=True)
            data_1d = pd.concat([data1d, now1d], ignore_index=True)
            data_5m = data_5m.drop_duplicates(subset=['datetime'], keep='first', ignore_index=True)
            data_1d = data_1d.drop_duplicates(subset=['datetime'], keep='first', ignore_index=True)
            start_date_time = data_5m['datetime'][0]  # 回测开始时间
            end_date_time = data_5m["datetime"].tail(1).values[0]
            data_5m.index = pd.to_datetime(data_5m["datetime"])
            data_1d.index = pd.to_datetime(data_1d["datetime"])
            stock_5m = self.stuck_data_5t(data_5m)
            stock_1d = self.stuck_data_5t(data_1d)
            stock_data_5t = StockPandasData(dataname=stock_5m)
            stock_data_1d = StockPandasData(dataname=stock_1d)
            self.cerebro.adddata(stock_data_5t, name=r['tdx_code'] + '_5m')
            self.cerebro.adddata(stock_data_1d, name=r['tdx_code'] + '_1d')
            self.cerebro.run()
            if rule == '15m':
                stock_data_15t = self.stuck_data_15t(data_5m)
            elif rule == '60m':
                stock_data_15t = self.stuck_data_15t(data_5m)
                stock_data_60t = self.stuck_data_60t(data_5m)
            elif rule == '1d':
                stock_data_15t = self.stuck_data_15t(data_5m)
                stock_data_60t = self.stuck_data_60t(data_5m)
                stock_data_1d = self.stuck_data_1d(data_5m)
            else:
                pass
        logging.info("执行任务完成，任务参数：%s" % rule)

    def add_job(self):
        # 添加任务并设置触发方式为3s一次
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='9', minute='36, 41, 51, 56', args=['5m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='9', minute='46', args=['15m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='11', minute='06, 11, 21, 26', args=['5m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='10', minute='01, 16, 46', args=['15m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='11', minute='01, 16', args=['15m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='10-11', minute='31', args=['1h'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='10, 13-15', minute='06, 16, 21, 26, 36, 41, 51, 56', args=['5m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='13-15', minute='16, 31, 46', args=['15m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='14-15', minute='01', args=['1h'])

    def scheduler_start(self):
        self.scheduler.start()

    def _read_data_5m(self, ths_code):
        stock_data = self.security.getStockBarsFromLocal(ths_code, '5m')
        start_date_time = stock_data['datetime'][0]  # 回测开始时间
        end_date_time = stock_data["datetime"].tail(1).values[0]
        return (stock_data, stock_data.shape[0], start_date_time, end_date_time)

    def _read_data_1d(self, ths_code):
        stock_data = self.security.getStockBarsFromLocal(ths_code, '1d')
        start_date_time = stock_data['datetime'][0]  # 回测开始时间
        end_date_time = stock_data["datetime"].tail(1).values[0]
        return (stock_data, stock_data.shape[0], start_date_time, end_date_time)

    def stuck_data_5t(self, data):
        stock_data_5t = data.join(self.analysis.ema(data, 21, 55, 144), how='outer')
        stock_data_5t = stock_data_5t.join(self.analysis.volume(data), how='outer')
        stock_data_5t = stock_data_5t.join(self.analysis.kdj(data), how='outer')
        stock_data_5t = stock_data_5t.join(self.analysis.cci(data), how='outer')
        stock_data_5t = stock_data_5t.join(self.analysis.rsi(data), how='outer')
        stock_data_5t = stock_data_5t.join(self.analysis.macd(data), how='outer')
        stock_data_5t = stock_data_5t.join(self.analysis.bollinger(data), how='outer')
        stock_data_5t = stock_data_5t.join(self.analysis.mfi(data), how='outer')
        #stock_data_5t = stock_data_5t.join(self.analysis.obv(data), how='outer')
        stock_data_5t = stock_data_5t.join(self.analysis.atr(data), how='outer')
        stock_data_5t = stock_data_5t.join(self.analysis.willr(data), how='outer')
        stock_data_5t = stock_data_5t.join(self.analysis.sar(data), how='outer')
        return stock_data_5t

    def stuck_data_15t(self, data):
        stock_data = self.security.lc5Resample(data, rule='15T')
        stock_data_15t = stock_data.join(self.analysis.ema(stock_data, 21, 55, 144), how='outer')
        stock_data_15t = stock_data_15t.join(self.analysis.volume(stock_data), how='outer')
        stock_data_15t = stock_data_15t.join(self.analysis.kdj(stock_data), how='outer')
        stock_data_15t = stock_data_15t.join(self.analysis.cci(stock_data), how='outer')
        stock_data_15t = stock_data_15t.join(self.analysis.rsi(stock_data), how='outer')
        stock_data_15t = stock_data_15t.join(self.analysis.macd(stock_data), how='outer')
        stock_data_15t = stock_data_15t.join(self.analysis.bollinger(stock_data), how='outer')
        stock_data_15t = stock_data_15t.join(self.analysis.mfi(stock_data), how='outer')
        #stock_data_15t = stock_data_15t.join(self.analysis.obv(stock_data), how='outer')
        stock_data_15t = stock_data_15t.join(self.analysis.atr(stock_data), how='outer')
        stock_data_15t = stock_data_15t.join(self.analysis.willr(stock_data), how='outer')
        stock_data_15t = stock_data_15t.join(self.analysis.sar(stock_data), how='outer')
        return stock_data_15t

    def stuck_data_60t(self, data):
        stock_data = self.security.lc5Resample(data, rule='60T')
        stock_data_60t = stock_data.join(self.analysis.ema(stock_data, 21, 55, 144), how='outer')
        stock_data_60t = stock_data_60t.join(self.analysis.volume(stock_data), how='outer')
        stock_data_60t = stock_data_60t.join(self.analysis.kdj(stock_data), how='outer')
        stock_data_60t = stock_data_60t.join(self.analysis.cci(stock_data), how='outer')
        stock_data_60t = stock_data_60t.join(self.analysis.rsi(stock_data), how='outer')
        stock_data_60t = stock_data_60t.join(self.analysis.macd(stock_data), how='outer')
        stock_data_60t = stock_data_60t.join(self.analysis.bollinger(stock_data), how='outer')
        stock_data_60t = stock_data_60t.join(self.analysis.mfi(stock_data), how='outer')
        #stock_data_60t = stock_data_60t.join(self.analysis.obv(stock_data), how='outer')
        stock_data_60t = stock_data_60t.join(self.analysis.atr(stock_data), how='outer')
        stock_data_60t = stock_data_60t.join(self.analysis.willr(stock_data), how='outer')
        stock_data_60t = stock_data_60t.join(self.analysis.sar(stock_data), how='outer')
        return stock_data_60t

    def stuck_data_1d(self, data):
        stock_data_1d = data.join(self.analysis.ema(data, 21, 55, 144), how='outer')
        stock_data_1d = stock_data_1d.join(self.analysis.kdj(data), how='outer')
        stock_data_1d = stock_data_1d.join(self.analysis.cci(data), how='outer')
        stock_data_1d = stock_data_1d.join(self.analysis.rsi(data), how='outer')
        stock_data_1d = stock_data_1d.join(self.analysis.macd(data), how='outer')
        stock_data_1d = stock_data_1d.join(self.analysis.bollinger(data), how='outer')
        stock_data_1d = stock_data_1d.join(self.analysis.mfi(data), how='outer')
        #stock_data_1d = stock_data_1d.join(self.analysis.obv(data), how='outer')
        stock_data_1d = stock_data_1d.join(self.analysis.atr(data), how='outer')
        stock_data_1d = stock_data_1d.join(self.analysis.willr(data), how='outer')
        stock_data_1d = stock_data_1d.join(self.analysis.sar(data), how='outer')
        return stock_data_1d

class TradeStrategy(bt.Strategy):
    def log(self, txt, dt=None):
        ''' 提供记录功能'''
        dt = dt or self.datas[0].datetime.datetime(0)
        print('%s, %s' % (dt.isoformat(sep=" "), txt))

    def __init__(self):
        #print('当前持仓量', self.getposition(self.data).size)
        #print('当前持仓成本', self.getposition(self.data).price)
        # self.getpositionbyname(name=None, broker=None)
        for i in range(0, len(self.datas)):
            print("--------- 打印 self.datas 中 Index:" + str(i) + "，数据集名称：" + self.datas[i]._name + " ----------")

        # Add a MovingAverageSimple indicator
        '''self.emaShort = bt.talib.EMA(self.data, timeperiod=21)
        self.emaMiddle = bt.talib.EMA(self.data, timeperiod=55)
        self.emaLong = bt.talib.EMA(self.data, timeperiod=144)'''

    def next(self):
        # 目前的策略就是简单显示下收盘价。
        #self.log('Close, %.2f，Volume： %.2f, emaS：%.2f，tr：%.2f，sar：%.2f，sarext：%.2f' % (self.dataclose[0], self.datavolume[0], self.datas[0].emaS[0], self.datas[0].tr[0], self.datas[0].sar[0], self.datas[0].sarext[0]))
        #self.log('Close： %.2f，Open： %.2f，Volume： %.2f，Ema(21): %.2f，Ema(55), %.2f，Ema(144), %.2f' % (self.dataclose[0], self.dataopen[0], self.datavolume[0], self.emaShort[0], self.emaMiddle[0], self.emaLong[0]))
        for i in range(0, len(self.datas)):
            data = self.getdatabyname(self.datas[i]._name)  # 根据名称返回数据集
            if(data.macd[0] >= 0 and data.diff[-1] < data.dea[-1] and data.diff[0] > data.dea[0] and data.diff[0] > 0):
                self.log(self.datas[i]._name + ', BUY, price: ' + str(data.close[0]) + '，macd：' + str(data.macd[0]) + '，diff：' + str(data.diff[0]) + '，dea：' + str(data.dea[0]))

class StockPandasData(bt.feeds.PandasData):
    lines = ('ratio', 'amount', 'volShort', 'volLong', 'emaS', 'emaM', 'emaL', 'K', 'D', 'J', 'rsi', 'cci', 'diff', 'dea', 'macd', 'bollLower', 'bollMiddle', 'bollUpper',
             'mfi', 'atr', 'trange', 'tr', 'willrL', 'willrS', 'sar', 'sarext')  # 要添加的线
    params = (
        ('nullvalue', 0.0),
        ('ratio', -1),
        ('amount', -1),
        ('volShort', -1),
        ('volLong', -1),
        ('emaS', -1),
        ('emaM', -1),
        ('emaL', -1),
        ('K', -1),
        ('D', -1),
        ('J', -1),
        ('rsi', -1),
        ('cci', -1),
        ('diff', -1),
        ('dea', -1),
        ('macd', -1),
        ('bollLower', -1),
        ('bollMiddle', -1),
        ('bollUpper', -1),
        ('mfi', -1),
        ('atr', -1),
        ('trange', -1),
        ('tr', -1),
        ('willrL', -1),
        ('willrS', -1),
        ('sar', -1),
        ('sarext', -1),
    )

if __name__ == '__main__':
    # 实例化一个调度器
    units = QUnits()
    setting = units.XML2ObjFromFile('./setting.xml').setting
    conn = units.dbConnect(setting.pgsql)
    cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
    cursor.execute("SELECT * FROM trading.owner_stock_view  WHERE account_id = trading.func_get_default_account('GadflyBSD', 'account.stock')->'data'->>'uuid' Limit 1")
    stock_schedule = StockSchedule(cursor.fetchall())
    stock_schedule.add_job()
    stock_schedule.scheduler_start()
    #stock_schedule = StockSchedule(cursor.fetchall())